/******************************************************************************
	This file is a part of the Snapshot Bioinformatics Project

	Copyright (C) 2012-2013 Jarom Schow

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
******************************************************************************/

#include "SqlServer.h"

#include <QtCore/QVariant>
#include <QtSql/QSqlDatabase>
#include <QtSql/QSqlQuery>
#include <QtSql/QSqlError>

#include "Core.h"
#include "MySqlEmbedded.h"
#include "Progress.h"
#include "SeqDatabase.h"
#include "SeqRange.h"
#include "SqlScripts.h"
#include "TaxonUpdater.h"
#include "Strings.h"

using namespace Bio;

///////////////////////////////////////////////////////////////////////////////
// SqlServerSettings
///////////////////////////////////////////////////////////////////////////////

/**
**/
SqlServerSettings::SqlServerSettings() : m_type( Unknown ), m_port( -1 )
{
}

/**
**/
SqlServerSettings::SqlServerSettings( const QString &configFile, const QString &dbname ) : m_type( MySQLEmbedded ),
	m_dbname( dbname ), m_configFile( configFile ), m_port( -1 )
{
}

/**
**/
SqlServerSettings::SqlServerSettings( ServerType type, const QString &host, int port, const QString &user, 
						const QString &pass, const QString &dbname ) : m_type( type ), 
	m_host( host ), m_port( port ), m_user( user ), m_password( pass ), m_dbname( dbname )
{
}

/**
**/
SqlServerSettings::SqlServerSettings( const SqlServerSettings &settings )
{
	*this = settings;
}

/**
**/
SqlServerSettings::~SqlServerSettings()
{
}

/**
**/
SqlServerSettings& SqlServerSettings::operator=( const SqlServerSettings &settings )
{
	m_port = settings.m_port;
	m_host = settings.m_host;
	m_user = settings.m_user;
	m_password = settings.m_password;
	m_dbname = settings.m_dbname;
	m_configFile = settings.m_configFile;
	m_type = settings.m_type;

	return *this;
}

/**
**/
bool SqlServerSettings::isValid()  const
{
	if( m_type == Unknown )
		return false;

	if( m_type == MySQLEmbedded )
		return !m_configFile.isEmpty();

	return	!(m_host.isEmpty() || m_user.isEmpty() || m_dbname.isEmpty());
}

/**
**/
void SqlServerSettings::setServerType( ServerType type )
{
	m_type = type;
}

/**
**/
SqlServerSettings::ServerType SqlServerSettings::getServerType() const
{
	return m_type;
}

/**
**/
void SqlServerSettings::setHostName( const QString &host )
{
	m_host = host;
}

/**
**/
QString SqlServerSettings::getHostName() const
{
	return m_host;
}

/**
**/
void SqlServerSettings::setPort( int port )
{
	m_port = port;
}

/**
**/
int SqlServerSettings::getPort() const
{
	return m_port;
}

/**
**/
void SqlServerSettings::setUserName( const QString &name )
{
	m_user = name;
}

/**
**/
QString SqlServerSettings::getUserName() const
{
	return m_user;
}

/**
**/
void SqlServerSettings::setPassword( const QString &pass )
{
	m_password = pass;
}

/**
**/
QString SqlServerSettings::getPassword() const
{
	return m_password;
}

/**
**/
void SqlServerSettings::setDatabaseName( const QString &name )
{
	m_dbname = name;
}

/**
**/
QString SqlServerSettings::getDatabaseName() const
{
	return m_dbname;
}

/**
**/
void SqlServerSettings::setConfigurationFile( const QString &configFile )
{
	m_configFile = configFile;
}

/**
**/
QString SqlServerSettings::getConfigurationFile() const
{
	return m_configFile;
}

///////////////////////////////////////////////////////////////////////////////
// SqlServer
///////////////////////////////////////////////////////////////////////////////

/**
**/
SqlServer::SqlServer() : m_connected( false ), m_inTransaction( false )
{
}

/**
**/
SqlServer::SqlServer( const SqlServerSettings &settings ) : m_settings( settings ),
	m_connected( false ), m_inTransaction( false )
{
	if( settings.m_type != SqlServerSettings::MySQLEmbedded )
		connect();
}

/**
**/
SqlServer::~SqlServer()
{
	close();
}

/**
**/
bool SqlServer::isConnected() const
{
	return m_connected;
}

/**
**/
bool SqlServer::connect()
{
	if( m_connected )
		return true;

	if( m_settings.m_type == SqlServerSettings::MySQLEmbedded ){
		m_db = MySqlEmbedded::addDatabase( m_settings.m_configFile, m_settings.m_dbname );
	}
	else {
		if( m_settings.m_type == SqlServerSettings::MySQL ){
			m_db = QSqlDatabase::addDatabase( "QMYSQL" );
		}
		else if( m_settings.m_type == SqlServerSettings::PostgreSQL ){
			m_db = QSqlDatabase::addDatabase( "QPSQL" );
		}
		else {
			setError( "Server type not set. Can not connect." );
			return false;
		}

		if( !m_db.isValid() ){
			setError( "Could not create driver for database." );
			return false;
		}

		m_db.setHostName( m_settings.m_host );
		m_db.setPort( m_settings.m_port );
		m_db.setUserName( m_settings.m_user );
		m_db.setPassword( m_settings.m_password );
		m_db.setDatabaseName( m_settings.m_dbname );
	}

	if( !m_db.open() ){
		setError( "Error connecting to database server: " + m_db.lastError().text() );
		return false;
	}

	m_connected = true;
	m_query = QSqlQuery( m_db );

	if( m_settings.m_type == SqlServerSettings::MySQLEmbedded ){
		// the embedded driver is not setting the database correctly - this is a hack until the driver gets fixed!
		QString	sql = "USE " + m_settings.m_dbname;
		if( !exec( sql ) )
			setError( "Error switching to database: " + m_db.lastError().text() );
	}

	return true;
}

/**
**/
bool SqlServer::connect( const SqlServerSettings &settings )
{
	m_settings = settings;

	return connect();
}

/**
**/
void SqlServer::close()
{
	if( !m_connected )
		return;

	m_db.close();
	m_connected = false;
}

/**
**/
bool SqlServer::create()
{
	if( m_connected )
		return false;

	if( m_settings.m_type == SqlServerSettings::MySQL || m_settings.m_type == SqlServerSettings::MySQLEmbedded ){
		return createMySQL();
	}
	else if( m_settings.m_type == SqlServerSettings::PostgreSQL ){
		return createPostgre();
	}

	setError( "Server type not set. Can not create database." );

	return false;
}

/**
**/
bool SqlServer::create( const SqlServerSettings &settings )
{
	m_settings = settings;

	return create();
}

/**
**/
bool SqlServer::createPostgre()
{
	Progress	progress( QString( "Creating database %1..." ).arg( m_settings.m_dbname ), -1 );

	m_db = QSqlDatabase::addDatabase( "QPSQL" );

	if( !m_db.isValid() ){
		setError( "Could not create driver for PostgreSQL database." );
		return false;
	}

	m_db.setHostName( m_settings.m_host );
	m_db.setPort( m_settings.m_port );
	m_db.setUserName( m_settings.m_user );
	m_db.setPassword( m_settings.m_password );

	if( !m_db.open() ){
		setError( "Error connecting to database server: " + m_db.lastError().text() );
		return false;
	}

	m_query = QSqlQuery( m_db );

	QString	sql = "CREATE DATABASE " + m_settings.m_dbname;

	if( !exec( sql ) ){
		setError( "Error creating database: " + m_db.lastError().text() );
		m_db.close();
		return false;
	}

	m_db.close();

	m_db.setDatabaseName( m_settings.m_dbname );

	if( !m_db.open() ){
		setError( "Error switching to database: " + m_db.lastError().text() );
		m_db.close();
		return false;
	}

	m_connected = true;

	sql = c_postgre_create_tables;
	if( !exec( sql ) ){
		setError( "Error building database tables: " + m_db.lastError().text() );
		m_db.close();
		return false;
	}

	progress.finish();

	return true;
}

/**
**/
bool SqlServer::createMySQL()
{
	Progress	progress( QString( "Creating database %1..." ).arg( m_settings.m_dbname ), -1 );

	if( m_settings.m_type == SqlServerSettings::MySQL ) {
		m_db = QSqlDatabase::addDatabase( "QMYSQL" );

		if( !m_db.isValid() ){
			setError( "Could not create driver for MySql database." );
			return false;
		}

		m_db.setHostName( m_settings.m_host );
		m_db.setPort( m_settings.m_port );
		m_db.setUserName( m_settings.m_user );
		m_db.setPassword( m_settings.m_password );

		if( !m_db.open() ){
			setError( "Error connecting to database server: " + m_db.lastError().text() );
			return false;
		}
	}
	else {
		m_db = MySqlEmbedded::addDatabase( m_settings.m_configFile );

		if( !m_db.isValid() ){
			setError( "Could not create driver for MySql database." );
			return false;
		}
	}

	m_query = QSqlQuery( m_db );

	QString	sql;

	sql = "CREATE DATABASE IF NOT EXISTS " + m_settings.m_dbname;

	if( !exec( sql ) ){
		setError( "Error creating database: " + m_db.lastError().text() );
		m_db.close();
		return false;
	}

	sql = "USE " + m_settings.m_dbname;
	if( !exec( sql ) ){
		setError( "Error switching to database: " + m_db.lastError().text() );
		m_db.close();
		return false;
	}

	m_connected = true;

	if( m_settings.m_type == SqlServerSettings::MySQL ) {
		sql = c_my_sql_create_tables;
		if( !exec( sql ) ){
			setError( "Error building database tables: " + m_db.lastError().text() );
			m_db.close();
			return false;
		}
	}
	else {
		const char	**script = c_my_sql_embedded_create_tables;

		while( *script ){
			sql = *script;
			if( !exec( sql ) ){
				setError( "Error building database tables: " + m_db.lastError().text() );
				m_db.close();
				return false;
			}
			script++;
		}
	}

	progress.finish();

	return true;
}

/**
**/
bool SqlServer::updateTaxonomy( const QString &tempDir )
{
	TaxonUpdater	updater;

	return updater.update( *this, tempDir );
}

/**
**/
const QSqlDatabase& SqlServer::getDatabase() const
{
	return m_db;
}

/**
**/
QStringList SqlServer::getSeqDatabaseNames() const
{
	QString		sql = QString( "SELECT name FROM biodatabase" );
	QSqlQuery	query( sql, m_db );
	QStringList	names;

	if( query.first() ) {
		do {
			names.append( query.value( 0 ).toString() );
		} while( query.next() );
	}

	return names;
}

/**
**/
SeqDatabase SqlServer::getSeqDatabase( const QString &name ) const
{
	QString		sql = "SELECT biodatabase_id FROM biodatabase WHERE name=" + esc( name );
	QSqlQuery	query( sql, m_db );

	if( !query.first() )
		return SeqDatabase(); // no database by that name

	int	bioID = query.value( 0 ).toInt();

	return SeqDatabase( name, bioID, const_cast<SqlServer*>( this ) );
}

/**
**/
bool SqlServer::createSeqDatabase( const QString &name, const QString &authority, const QString &description )
{
	QString	sql = "INSERT INTO biodatabase (name, authority, description) VALUES (";
		
	sql += esc( name ) + ", ";
	sql += esc( authority ) + ", ";
	sql += esc( description ) + ")";

	if( !exec( sql ) )
		return false;

	return true;
}

/**
**/
int SqlServer::getOntologyId( const QString &name )
{
	if( m_ontologyMap.contains( name ) )
		return m_ontologyMap[name];

	QString	sql = "SELECT ontology_id FROM ontology WHERE name = " + esc( name );

	if( !exec( sql ) )
		return -1;

	int	id = -1;

	if( m_query.first() )
		id = m_query.value( 0 ).toInt();
	else
		id = addOntology( name, QString::null );

	if( id >= 0 )
		m_ontologyMap[name] = id;
	
	return id;
}

/**
**/
int SqlServer::getTermId( const QString &term, const QString &ontology )
{
	int	ontologyId = getOntologyId( ontology );

	if( ontologyId < 0 )
		return -1;

	return getTermId( term, ontologyId );
}

/**
**/
int SqlServer::getTermId( const QString &term, int ontologyId )
{
	QHash<QString, int>	&termMap = m_termMaps[ontologyId];

	if( termMap.contains( term ) )
		return termMap[term];

	QString		sql = "SELECT term_id FROM term WHERE name = " + esc( term );

	if( ontologyId >= 0 )
		sql += " AND ontology_id = " + QString::number( ontologyId );

	if( !exec( sql ) )
		return -1;

	int	id = -1;

	if( m_query.first() )
		id =  m_query.value( 0 ).toInt();
	else
		id = addTerm( term, ontologyId, QString::null, QString::null );

	if( id >= 0 )
		termMap[term] = id;

	return id;
}

/**
**/
SeqRange SqlServer::getFeatureRange( int seqFeatureId )
{
	SeqRange	range;
	QString		sql = "SELECT start_pos, end_pos, strand FROM location WHERE seqfeature_id=" +
						QString::number( seqFeatureId) + " ORDER BY rank";

	if( !exec( sql ) )
		return range;

	if( !m_query.first() )
		return range;

	do {
		int						strandVal = m_query.value( 2 ).toInt();
		SeqLocation::Strand	strand = SeqLocation::Unknown;

		if( strandVal == 1 )
			strand = SeqLocation::Positive;
		else if( strandVal == -1 )
			strand = SeqLocation::Negative;

		range.addLocation( SeqLocation( m_query.value( 0 ).toInt(), m_query.value( 1 ).toInt(), strand ) );
	} while( m_query.next() );

	return range;
}

/**
**/
bool SqlServer::exec( const QString &sql )
{
	if( !m_query.exec( sql ) ){
		QString	sqlError = "SQL error: " + m_query.lastError().text();

		setError( sqlError );
		Core::log( "SQL statement: " + sql );
		
		return false;
	}

	return true;
}

/**
**/
bool SqlServer::startTransaction()
{
	if( m_inTransaction ){
		setError( "SqlServer::startTransaction(): Already in transaction, call commit() or rollback()." );
		return false;
	}

	if( !m_db.transaction() ){
		QString	sqlError = "Unable to start transaction: " + m_db.lastError().text();

		setError( sqlError );
		return false;
	}

	m_inTransaction = true;

	return true;
}

/**
**/
bool SqlServer::commit()
{
	if( !m_inTransaction ){
		setError( "SqlServer::commit(): No active transaction." );
		return false;
	}

	m_inTransaction = false;

	if( !m_db.commit() ){
		QString	sqlError = "Error trying to commit transaction: " + m_db.lastError().text();

		setError( sqlError );
		return false;
	}

	return true;
}

/**
**/
bool SqlServer::rollback()
{
	if( !m_inTransaction ){
		setError( "SqlServer::rollback(): No active transaction." );
		return false;
	}

	m_inTransaction = false;

	if( !m_db.rollback() ){
		QString	sqlError = "Error trying to rollback transaction: " + m_db.lastError().text();

		setError( sqlError );
		return false;
	}

	return true;
}

/**
**/
QSqlQuery SqlServer::getQuery() const
{
	return m_query;
}

/**
**/
QString SqlServer::getLastError() const
{
	return m_lastError;
}

/**
**/
QString SqlServer::esc( const QString &str, bool emptyAsNull ) const
{
	if( emptyAsNull && str.isEmpty() )
		return "NULL";

	QString	ret;
	bool	escaped = false;

	ret.reserve( str.length() + 2 );

	ret += '\'';

	for( QString::const_iterator it = str.begin(); it != str.end(); it++ ){
		if (*it == '\"'){
			ret += "\\\"";
			escaped = true;
		}
		else if (*it == '\''){
			ret += "\\'";
			escaped = true;
		}
		else if (*it == '\\'){
			ret += "\\\\";
			escaped = true;
		}
		else
			ret += *it;
	}

	ret += '\'';

	if( escaped && m_settings.m_type == SqlServerSettings::PostgreSQL )
		return "E" + ret;

	return ret;
}

/**
**/
SeqRecord SqlServer::getRecord( int bioentryId, int seqDatabaseId )
{
	return getRecord( "WHERE bioentry_id = " + QString::number( bioentryId ) );
}

/**
**/
SeqRecord SqlServer::getRecordByAccession( const QString &accession, int seqDatabaseId )
{
	return getRecord( "WHERE accession = " + esc( accession ) );
}

/**
**/
QList<SeqRecord> SqlServer::getRecordsForTaxon( int taxonId, int seqDatabaseId )
{
	QList<SeqRecord>	records;

	QString sql = "SELECT bioentry_id, taxon_id, name, accession, identifier, "
					"division, description, version FROM bioentry WHERE taxon_id = " + 
					QString::number( taxonId );

	if( !exec( sql ) )
		return records;

	if( !m_query.first() )
		return records;

	do {
		records.append( makeRecord( m_query ) );
	} while( m_query.next() );

	return records;
}

/**
**/
Sequence* SqlServer::getSequence( int bioentryId )
{
	QString sql = "SELECT alphabet, length, seq FROM biosequence WHERE bioentry_id = " + QString::number( bioentryId );

	if( !exec( sql ) )
		return NULL;

	if( !m_query.first() )
		return NULL;

	Sequence	*seq = new Sequence;
	int			length;
	QString		alphabet;

	alphabet = m_query.value( 0 ).toString();
	length = m_query.value( 1 ).toInt();
	seq->setSequence( m_query.value( 2 ).toByteArray() );

	if( alphabet == c_alphaDNA )
		seq->setAlphabet( Sequence::DNA );
	else if( alphabet == c_alphaRNA )
		seq->setAlphabet( Sequence::RNA );
	else if( alphabet == c_alphaProtein )
		seq->setAlphabet( Sequence::Protein );
	else
		seq->setAlphabet( Sequence::Unknown );

	return seq;
}

/**
**/
bool SqlServer::createUser( const QString &userName, const QString &password, bool grantAll )
{
	QString	sql = "CREATE USER " + esc( userName );

	if( m_settings.m_type == SqlServerSettings::PostgreSQL ){
		if( !password.isEmpty() )
			sql += " WITH PASSWORD " + esc( password );
	}
	else if( m_settings.m_type == SqlServerSettings::MySQL ){
		if( !password.isEmpty() )
			sql += " IDENTIFIED BY " + esc( password );
	}
	else
		return false;

	if( !exec( sql ) )
		return false;

	if( grantAll )
		sql = "GRANT ALL ON ";
	else
		sql = "GRANT SELECT ON ";

	if( m_settings.m_type == SqlServerSettings::PostgreSQL ){
		sql += " DATABASE " + m_settings.m_dbname;
	}
	else if( m_settings.m_type == SqlServerSettings::MySQL ){
		sql += m_settings.m_dbname + ".*";
	}
	else
		return false;

	sql += " TO " + esc( userName );

	return exec( sql );
}

/**
**/
int SqlServer::addOntology( const QString &name, const QString &definition )
{
	QString	sql = "INSERT INTO ontology(name, definition) VALUES (";
	
	sql += esc( name ) + ", " + esc( definition, true ) + ")";

	if( !exec( sql ) )
		return -1;

	QVariant	var = m_query.lastInsertId();
	
	if( var.isValid() )
		return var.toInt();

	sql = "SELECT ontology_id FROM ontology WHERE name = " + esc( name );

	if( !exec( sql ) )
		return -1;

	if( m_query.first() )
		return m_query.value( 0 ).toInt();

	return -1;
}

/**
**/
int SqlServer::addTerm( const QString &term, int ontologyId, const QString &definition, const QString &identifier )
{
	QString	sql = "INSERT INTO term (name, definition, identifier, ontology_id) VALUES (";
		
	sql += esc( term ) + ", ";
	sql += esc( definition, true ) + ", ";
	sql += esc( identifier, true ) + ", ";
	sql += QString::number( ontologyId ) + ")";

	if( !exec( sql ) )
		return -1;

	QVariant	var = m_query.lastInsertId();
	
	if( var.isValid() )
		return var.toInt();

	sql = "SELECT term_id FROM term WHERE name = " + esc( term ) + " AND ontology_id = " + QString::number( ontologyId );

	if( !exec( sql ) )
		return -1;

	if( m_query.first() )
		return m_query.value( 0 ).toInt();

	return -1;
}

/**
**/
void SqlServer::setError( const QString &msg )
{
	m_lastError = msg;
	Core::error( msg );
}

/**
**/
SeqRecord SqlServer::getRecord( const QString &condition )
{
	QString sql( "SELECT bioentry_id, taxon_id, name, accession, identifier, "
				"division, description, version FROM bioentry" );

	if( !condition.isEmpty() )
		sql += " " + condition;
	
	if( !exec( sql ) )
		return SeqRecord();

	if( !m_query.first() )
		return SeqRecord();

	return makeRecord( m_query );
}

/**
**/
SeqRecord SqlServer::makeRecord( const QSqlQuery &result )
{
	int			version = 0;
	int			bioId = result.value( 0 ).toInt();
	QString		identifier, division;
	SeqRecord	record( bioId, this );

	record.setTaxonId( result.value( 1 ).toInt() );
	record.setName( result.value( 2 ).toString() );
	record.setId( result.value( 3 ).toString() );
	identifier = result.value( 4 ).toString();
	division = result.value( 5 ).toString();
	record.setDescription( result.value( 6 ).toString() );
	version = result.value( 7 ).toInt();

	if( !identifier.isEmpty() )
		record.setAnnotation( c_annGenBankId, identifier );

	if( !division.isEmpty() )
		record.setAnnotation( c_annDataFileDivision, division );

	if( version > 0 )
		record.setAnnotation( c_annSequenceVersion, QString::number( version ) );
/*
	Sequence	*seq = getSequence( record.getBioentryId() );

	if( seq )
		record.setSequence( seq );

	getDbXRefs( record );
	getFeatures( record );
	getAnnotations( record );
*/
	return record;
}
